package modelarts

import (
	"fmt"
	"strings"
)

const (
	// 游客
	GuestUserType UserType = "guest"
	// 游客
	OwnerUserType UserType = "owner"
	// 推理
	Inference ResourceType = "inference"
	// 日志可视化
	Logsvisual ResourceType = "logsvisual"
	// 用户自定义日志可视化;
	LogsvisualCustomize ResourceType = "logsvisual_customize"
)

type UserType string
type ResourceType string

type Metrics struct {
	LearningRate string `json:"learning_rate"`
	BatchSize    string `json:"batch_size"`
	Momentum     string `json:"momentum"`
}

type Aim struct {
	modelArtsProject *ModelArtsProject
	// 羲和的用户名
	userName string
	// 羲和的项目名
	projectName string
	// 训练任务的id
	jobID      string
	apiVersion string
	namespace  string
}

// 初始化一个Aim对象
func NewAim(username, projectname, jobID string) (aim *Aim) {
	aim = new(Aim)
	aim.userName = username
	aim.projectName = projectname
	aim.jobID = jobID
	aim.namespace = "mindspore-xihe"
	aim.apiVersion = fmt.Sprintf("%s/%s", "cs.opensourceways.com", "v1alpha1")
	return
}

// resource_type: inference:推理; logsvisual:日志可视化; logsvisual_customize:用户自定义日志可视化;
// user_type: guest:游客; owner: 拥有者;
func (aim *Aim) MakeResource(resourceType ResourceType, fileString, logPath, dbPath string, metrics *Metrics, userType UserType) string {
	resourceName := aim.makeResourceName(resourceType, userType)
	crdConfig := ""
	switch resourceType {
	case Inference:
		crdConfig = aim.makeGradioResourceTemp(resourceName, fileString)
	case Logsvisual:
		crdConfig = aim.makeLogsvisualResourceTemp(resourceName, logPath, metrics)
	case LogsvisualCustomize:
		crdConfig = aim.makeLogsvisualCustResourceTemp(resourceName, dbPath)
	}
	return crdConfig
}

func (aim *Aim) makeResourceName(resourceType ResourceType, userType UserType) (resourceName string) {
	switch resourceType {
	case Inference:
		resourceName = fmt.Sprintf("%s-%s-inference", aim.userName, aim.projectName)
	case Logsvisual:
		resourceName = fmt.Sprintf("%s-%s-logsvisual", aim.userName, aim.projectName)
	case LogsvisualCustomize:
		resourceName = fmt.Sprintf("%s-%s-custlogsvisual", aim.userName, aim.projectName)
	}
	resourceName = strings.ToLower(resourceName)
	resourceName = strings.ReplaceAll(resourceName, "_", "-")
	if len(userType) > 0 {
		resourceName = fmt.Sprintf("%s-%s", resourceName, userType)
	}
	return
}

func (aim *Aim) makeGradioResourceTemp(resourceName, fileString string) string {
	cmdStr := fmt.Sprintf(`python3 obs_folder_download.py --source_dir='xihe-obj/projects/%s/%s/inference/'
	 --source_files='%s' --dest='%s' --obs-ak=%s --obs-sk=%s --obs-bucketname=%s --obs-endpoint=%s\n cd /workspace/content\n 
	 pip install --upgrade -i https://pypi.tuna.tsinghua.edu.cn/simple pip\npip install -r requirements.txt -i
	  https://pypi.tuna.tsinghua.edu.cn/simple\n python3 app.py")`, aim.userName, aim.projectName, fileString, "/workspace/content/",
		aim.modelArtsProject.Manager.IamAK, aim.modelArtsProject.Manager.IamSK, aim.modelArtsProject.Manager.ObsBucketName, aim.modelArtsProject.obsEndPointUrl)
	return `{
		"apiVersion": "` + aim.apiVersion + `",
		"kind": "CodeServer",
		"metadata": {
			"name": ` + resourceName + `,
			"namespace": ` + aim.namespace + `
		},
		"spec": {
			"runtime": "generic",
			"subdomain": ` + resourceName + `,
			"image": "swr.cn-north-4.myhuaweicloud.com/opensourceway/xihe/gradio:0af9e11dbd10e10aa25ad22716e48adbbaf670fb",
			"storageSize": "1Gi",
			"storageName": "emptyDir",
			"inactiveAfterSeconds": 0,
			"recycleAfterSeconds": 60*60*24,
			"restartPolicy": "Never",
			"resources": {
			"requests": {
				"cpu": "0.5",
				"memory": "512Mi"
			}
			},
			"connectProbe": "/",
			"workspaceLocation": "/workspace",
			"envs": [
			{
				"name": "GRADIO_SERVER_PORT",
				"value": "8080"
			},
			{
				"name": "GRADIO_SERVER_NAME",
				"value": "0.0.0.0"
			}
			],
			"command": [ 
			"/bin/bash",
			"-c",
			 ` + cmdStr + `
			]
		}
	}`
}
func (aim *Aim) makeLogsvisualResourceTemp(resourceName, logPath string, metrics *Metrics) string {
	cmdStr := fmt.Sprintf(`"export PATH=/root/.local/bin:$PATH\n python3 aimTrace.py --log_path %s --repo /workspace 
	--learning_rate_scope '%s' --batch_size_scope '%s' --momentum_scope '%s' --obs_ak %s --obs_sk %s 
	--obs_bucketname %s --obs_endpoint %s\n aim up --repo /workspace --host 0.0.0.0 --port 8080 
	--workers 2\n"`, logPath, metrics.LearningRate, metrics.BatchSize, metrics.Momentum, aim.modelArtsProject.Manager.IamAK,
		aim.modelArtsProject.Manager.IamSK, aim.modelArtsProject.Manager.ObsBucketName, aim.modelArtsProject.obsEndPointUrl)
	return `{
		"apiVersion": "` + aim.apiVersion + `",
		"kind": "CodeServer",
		"metadata": {
			"name":` + resourceName + `,
			"namespace":` + aim.namespace + `
		},
		"spec": {
			"runtime": "generic",
			"subdomain": ` + resourceName + `,
			"image": "swr.cn-north-4.myhuaweicloud.com/opensourceway/xihe/xihe-aim-container:5dd5861416dccd7ada32b34ab8b83e743a245ce6",
			"storageSize": "1Gi",
			"storageName": "emptyDir",
			"inactiveAfterSeconds": 600,
			"recycleAfterSeconds": 1200,
			"restartPolicy": "Never",
			"resources": {
			"requests": {
				"cpu": "0.5",
				"memory": "512Mi"
			}
			},
			"connectProbe": "/",
			"workspaceLocation": "/workspace",
			"command": [
			"/bin/bash",
			"-c",
			# "tail -f README.md"
			` + cmdStr + ` 
			]
		}
}`
}
func (aim *Aim) makeLogsvisualCustResourceTemp(resourceName, dbPath string) string {
	cmdStr := fmt.Sprintf(`"export PATH=/root/.local/bin:$PATH\n python3 obs_folder_download.py --source %s  --dest /workspace/ 
	--obs-ak %s --obs-sk %s --obs-bucketname %s --obs-endpoint %s\n ls /workspace/\n aim up --repo /workspace/ 
	--host 0.0.0.0 --port 8080 --workers 2 "`, dbPath, aim.modelArtsProject.Manager.IamAK, aim.modelArtsProject.Manager.IamSK,
		aim.modelArtsProject.Manager.ObsBucketName, aim.modelArtsProject.obsEndPointUrl)
	return `{
		"apiVersion": "` + aim.apiVersion + `",
		"kind": "CodeServer",
		"metadata": {
			"name": ` + resourceName + `,
			"namespace": ` + aim.namespace + `
		},
		"spec": {
			"runtime": "generic",
			"subdomain":` + resourceName + `,
			"image": "swr.cn-north-4.myhuaweicloud.com/opensourceway/xihe/xihe-aim-container:5dd5861416dccd7ada32b34ab8b83e743a245ce6",
			"storageSize": "1Gi",
			"storageName": "emptyDir",
			"inactiveAfterSeconds": 600,
			"recycleAfterSeconds": 1200,
			"restartPolicy": "Never",
			"resources": {
			"requests": {
				"cpu": "0.5",
				"memory": "512Mi"
			}
			},
			"connectProbe": "/",
			"workspaceLocation": "/workspace",
			"command": [
			"/bin/bash",
			"-c",
			` + cmdStr + `  
			]
		}
}`
}
